import numpy as np
import pandas as pd


# 接收一个history_candles,增加新列['signal'] 1 为买入信号,2为卖出信号
def macdFilterstartegy(data_df, short=12, long=26, mid=9):
    history_candles = data_df.copy()
    history_candles.sort_values(by="ts", ascending=True, inplace=True)
    short_ema = history_candles["c"].ewm(span=short).mean()
    long_ema = history_candles["c"].ewm(span=long).mean()
    history_candles.loc[:, 'DIFF'] = short_ema - long_ema
    history_candles.loc[:, 'DEA'] = history_candles['DIFF'].ewm(span=mid).mean()
    history_candles.loc[:, 'MACD'] = 2 * (history_candles['DIFF'] - history_candles['DEA'])
    history_candles.loc[:, 'signal'] = 0  # 初始化signal全为0
    history_candles.reset_index(drop=True, inplace=True)  # 此时0对应的是最老的candle
    # 此时得到的history_candlse索引是倒序的
    direction = None
    for i in range(len(history_candles)):
        # 最开始的3个不要了
        if i <= 2:
            continue
        # 首先判断是否存在上穿或者下穿
        if history_candles.loc[i - 1, 'DEA'] * history_candles.loc[i - 2, 'DEA'] < 0:
            if history_candles.loc[i - 1, 'DEA'] > 0:
                direction = 'up'
                highest_price = max(history_candles.loc[i - 1, 'h'], history_candles.loc[i - 2, 'h'],
                                    history_candles.loc[i - 3, 'h'])
            else:
                direction = 'down'
                lowest_price = min(history_candles.loc[i - 1, 'l'], history_candles.loc[i - 2, 'l'],
                                   history_candles.loc[i - 3, 'l'])
            # 直接跳过了
            continue
        if direction == 'up':
            highest_price = max(highest_price, history_candles.loc[i - 2, 'h'])
            # 上穿成功
            if history_candles.loc[i - 1, 'c'] >= highest_price:
                history_candles.loc[i - 1, 'signal'] = 1
        if direction == 'down':
            lowest_price = min(lowest_price, history_candles.loc[i - 2, 'l'])
            if history_candles.loc[i - 1, 'c'] <= lowest_price:
                history_candles.loc[i - 1, 'signal'] = 2

    if 'Unnamed: 0' in history_candles:
        history_candles.drop(columns=['Unnamed: 0', 'MACD', 'DIFF', 'DEA', '1', '2', '3', 'complete'], inplace=True)
    else:
        history_candles.drop(columns=['MACD', 'DIFF', 'DEA', '1', '2', '3', 'complete'], inplace=True)
    return history_candles


if __name__ == '__main__':
    data_df = pd.read_csv('../data_info/bitcoin.csv')
    result = macdFilterstartegy(data_df)
    result.to_csv("./macdtest2.csv", index=False)
